#include "aes.h"

#include "debug.h"

static FL_ErrorStatus AES_GroupWriteAndRead_128BIT_IVR(uint8_t *dataIn, uint8_t *dataOut);
static FL_ErrorStatus AES_GroupWriteAndRead_128BIT(uint8_t *dataIn, uint8_t *dataOut);
static void           AES_DOUT_GroupRead_128BIT(uint8_t *dataOut);
static void           AES_DIN_GroupWrite_128BIT(uint8_t *dataIn);
static void           AES_IVR_ReadEx(uint8_t *ivrOut);
static void           AES_IVR_WriteEx(uint8_t *ivrIn);
static void           AES_KEY_WriteEx(uint8_t *keyIn, uint8_t len);
static void           AES_H_WriteEx(uint8_t *hxIn);

/**
 * @brief 初始化AES加密模块
 *
 * @param key_len           密钥长度
 * @param cipher_mode       加密模式
 * @param operation_mode    操作模式
 * @param data_type         数据类型
 */
void AES_Init(uint32_t key_len, uint32_t cipher_mode, uint32_t operation_mode, uint32_t data_type)
{
    FL_AES_InitTypeDef aes_init;

    /* 重置AES模块 */
    FL_AES_DeInit();

    /* 设置 AES 配置 */
    FL_AES_StructInit(&aes_init);
    aes_init.keyLength     = key_len;
    aes_init.cipherMode    = cipher_mode;
    aes_init.operationMode = operation_mode;
    aes_init.dataType      = data_type;

    /* 初始化AES */
    FL_AES_Init(AES, &aes_init);
}

/**
 * @brief  AES加解密处理流程
 * @note
 * DataOut指针空间由调用者申请，空间大小必须是16字节的整数倍（因为即使是不足16字节的数据进行加密，AES处理不足16字节会进行补零操作，然后再进行加解密）
 * @param    KeyIn 密钥指针
 * @param    KeyLen 密钥长度
 * @param    IVRIn 向量指针
 * @param    DataIn 加密输入数据指针
 * @param    DataInLen 加密数据输入长度
 * @param    DataOut 解密输出数据指针
 * @param    DataOutLen 解密输出长度
 * @retval 返回错误状态，可能值：
 *         -FL_PASS 测试成功
 *         -FL_FAIL 测试失败
 */
FL_ErrorStatus AES_EncryptionAndDecryption(uint8_t *key_in, uint8_t key_len, uint8_t *ivr_in,
                                           uint8_t *data_in, uint32_t data_in_len, uint8_t *data_out,
                                           uint32_t *data_out_len)
{
    uint32_t i;
    uint32_t result = 0;
    uint8_t  tempBuf[16];

    /* 设置密钥 */
    if (FL_AES_GetCipherMode(AES) != FL_AES_CIPHER_MULTH) {
        AES_KEY_WriteEx(key_in, key_len);
    }

    else {
        AES_H_WriteEx(key_in);
    }

    /* CTR和CBC模式下写AES_IVRx寄存器 */
    if ((ivr_in != NULL) && (FL_AES_GetCipherMode(AES) != FL_AES_CIPHER_ECB)) {
        AES_IVR_WriteEx(ivr_in);
    }

    /* 启动AES */
    FL_AES_Enable(AES);

    for (i = 0; i < data_in_len / 16; i++) {
        if (FL_AES_GetCipherMode(AES) == FL_AES_CIPHER_MULTH) {
            result = AES_GroupWriteAndRead_128BIT_IVR(data_in + 16 * i, data_out + 16 * i);

            if (result == FL_FAIL) {
                return FL_FAIL;
            }
        }

        else {
            result = AES_GroupWriteAndRead_128BIT(data_in + 16 * i, data_out + 16 * i);

            if (result == FL_FAIL) {
                return FL_FAIL;
            }
        }
    }

    if (data_in_len % 16) {
        memset(tempBuf, 0, sizeof(tempBuf));
        memcpy(tempBuf, data_in + i * 16, data_in_len - i * 16);

        if (FL_AES_GetOperationMode(AES) == FL_AES_CIPHER_MULTH) {
            result = AES_GroupWriteAndRead_128BIT_IVR(tempBuf, data_out + i * 16);

            if (result == FL_FAIL) {
                return FL_FAIL;
            }
        }

        else {
            result = AES_GroupWriteAndRead_128BIT(tempBuf, data_out + i * 16);

            if (result == FL_FAIL) {
                return FL_FAIL;
            }
        }
    }

    if (FL_AES_GetCipherMode(AES) == FL_AES_CIPHER_CTR) {
        *data_out_len = data_in_len;
    }

    else {
        *data_out_len = ((data_in_len + 15) / 16) * 16;
    }

    FL_AES_Disable(AES);
    return FL_PASS;
}

FL_ErrorStatus Test_AES_CTR_Decrypt(void)
{
    // 密文数据
    uint8_t dataInBuf[41] = { 0xDC, 0xCA, 0xB2, 0x5E, 0x5C, 0xCA, 0xFA, 0x88, 0xC1, 0x8B, 0xFF,
                              0x1B, 0x3D, 0xD8, 0x4A, 0x97, 0xDC, 0xCA, 0xB2, 0x5E, 0x5C, 0xCA,
                              0xFA, 0x88, 0xC1, 0x8B, 0xFF, 0x1B, 0x3D, 0xD8, 0x4A, 0x97, 0xDC,
                              0xCA, 0xB2, 0x5E, 0x5C, 0xCA, 0xFA, 0x88, 0xC1 };

    // 相同的密钥
    uint8_t keyBuf[16] = { 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
                           0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x31 };

    // 相同的初始向量
    uint8_t ivrBuf[16] = { 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30,
                           0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x32 };

    uint8_t  dataOut[50] = { 0x00 };
    uint32_t dataOutLen;

    // 期望的解密结果
    // const uint8_t scrResultBuf[16] = { 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37, 0x38,
    //                                    0x39, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };

    FL_AES_InitTypeDef AES_InitStructer;
    uint32_t           result;

    // 注意：CTR模式解密也使用加密操作模式
    FL_AES_Disable(AES);
    AES_InitStructer.keyLength     = FL_AES_KEY_LENGTH_128B;
    AES_InitStructer.cipherMode    = FL_AES_CIPHER_CTR;
    AES_InitStructer.operationMode = FL_AES_OPERATION_MODE_ENCRYPTION; // 仍然是加密模式！
    AES_InitStructer.dataType      = FL_AES_DATA_TYPE_32B;
    FL_AES_Init(AES, &AES_InitStructer);

    // 执行CTR解密（实际上是相同的加密操作）
    AES_EncryptionAndDecryption(keyBuf, 16, ivrBuf, dataInBuf, 41, dataOut, &dataOutLen);

    uint8_t  decode_out[50] = { 0 };
    uint32_t decode_out_len;

    AES_EncryptionAndDecryption(keyBuf, 16, ivrBuf, dataOut, dataOutLen, decode_out, &decode_out_len);
    for (uint8_t i = 0; i < decode_out_len; i++) {
        LOG("decode_out[%d]: %02X\n", i, decode_out[i]);
    }

    LOG("dataOutLen: %d\n", dataOutLen);

    // 验证结果
    // result = memcmp(dataOut, scrResultBuf, 16);

    // if (result != 0) {
    //     return FL_FAIL;
    // }
    return FL_PASS;
}

// /**
//  * @brief  数据字节交换
//  * @param  pWord: 数据指针
//  * @param  length 数据长度
//  * @retval None
//  */
// static void ByteExChange(uint32_t *pWord, uint32_t length)
// {
//     int32_t   i     = 0;
//     int32_t   value = 0;
//     uint32_t *pRam  = pWord;

//     for (i = 0; i < length; i++) {
//         value = (*pRam << 24) + (((*pRam & 0xFF00) >> 8) << 16) + (((*pRam & 0xFF0000) >> 16) << 8) +
//                 (*pRam >> 24);
//         *pRam++ = value;
//     }
// }

/**
 * @brief  设置MultH运算128bit输入H参数(h0存放密钥最低32bit)
 * @param  H参数 16字节长度的密钥
 * @param  长度  16
 * @retval None
 */
static void AES_H_WriteEx(uint8_t *hxIn)
{
    uint8_t       i;
    __IO uint32_t temp32;

    for (i = 0; i < 16; i = i + 4) {
        temp32 = (hxIn[0 + i] << 24) | (hxIn[1 + i] << 16) | (hxIn[2 + i] << 8) | hxIn[3 + i];
        FL_AES_WriteHParams(AES, (16 - i) / 4 - 1, temp32);
    }
}

// /**
//  * @brief  读取MultH运算128bit输入H参数(h0存放密钥最低32bit)
//  * @param  H参数 16字节长度的密钥
//  * @param  长度  16
//  * @retval None
//  */
// static void AES_H_ReadEx(uint8_t *hxOut)
// {
//     uint8_t       i;
//     __IO uint32_t temp32;

//     for (i = 0; i < 16; i = i + 4) {
//         temp32       = FL_AES_ReadHParams(AES, (16 - i) / 4 - 1);
//         hxOut[0 + i] = temp32 >> 24;
//         hxOut[1 + i] = temp32 >> 16;
//         hxOut[2 + i] = temp32 >> 8;
//         hxOut[3 + i] = temp32;
//     }
// }

/**
 * @brief  设置加解密密钥(key0存放密钥最低32bit)
 * @param  密钥  16/24/32字节长度的密钥
 * @param  长度  16/24/32
 * @retval None
 */
static void AES_KEY_WriteEx(uint8_t *keyIn, uint8_t len)
{
    uint8_t       i;
    __IO uint32_t temp32;

    if ((len != 16) && (len != 24) && (len != 32)) {
        return;
    }

    for (i = 0; i < len; i = i + 4) {
        temp32 = (keyIn[0 + i] << 24) | (keyIn[1 + i] << 16) | (keyIn[2 + i] << 8) | keyIn[3 + i];
        FL_AES_WriteKeys(AES, (len - i) / 4 - 1, temp32);
    }
}

// /**
//  * @brief  读取加解密密钥，第一个字存放密钥最低32bit
//  * @param  密钥  16/24/32字节长度的密钥
//  * @param  长度  16/24/32
//  * @retval None
//  */
// static void AES_KEY_ReadEx(uint8_t *keyOut, uint8_t len)
// {
//     uint8_t       i;
//     __IO uint32_t temp32;

//     if ((len != 16) && (len != 24) && (len != 32)) {
//         return;
//     }

//     for (i = 0; i < len; i = i + 4) {
//         temp32        = FL_AES_ReadKeys(AES, (len - i) / 4 - 1);
//         keyOut[0 + i] = temp32 >> 24;
//         keyOut[1 + i] = temp32 >> 16;
//         keyOut[2 + i] = temp32 >> 8;
//         keyOut[3 + i] = temp32;
//     }
// }

/**
 * @brief  设置初始向量
 * @param  初始向量
 * @retval None
 */
static void AES_IVR_WriteEx(uint8_t *ivrIn)
{
    uint8_t       i;
    __IO uint32_t temp32;

    for (i = 0; i < 4 * 4; i = i + 4) {
        temp32 = (ivrIn[0 + i] << 24) | (ivrIn[1 + i] << 16) | (ivrIn[2 + i] << 8) | ivrIn[3 + i];
        FL_AES_WriteIVR(AES, (16 - i) / 4 - 1, temp32);
    }
}

/**
 * @brief  读取初始向量
 * @param  初始向量
 * @retval None
 */
static void AES_IVR_ReadEx(uint8_t *ivrOut)
{
    uint8_t       i;
    __IO uint32_t temp32;

    for (i = 0; i < 4 * 4; i = i + 4) {
        temp32        = FL_AES_ReadIVR(AES, (4 * 4 - i) / 4 - 1);
        ivrOut[0 + i] = temp32 >> 24;
        ivrOut[1 + i] = temp32 >> 16;
        ivrOut[2 + i] = temp32 >> 8;
        ivrOut[3 + i] = temp32;
    }
}

/**
 * @brief  设置加解密数据，16字节（128bit）分组数据输入
 * @param  初始向量
 * @retval None
 */
static void AES_DIN_GroupWrite_128BIT(uint8_t *dataIn)
{
    uint8_t       i;
    __IO uint32_t temp32;

    for (i = 0; i < 16; i = i + 4) {
        temp32 = (dataIn[0 + i] << 24) | (dataIn[1 + i] << 16) | (dataIn[2 + i] << 8) | dataIn[3 + i];
        FL_AES_WriteInputData(AES, temp32);
    }
}

/**
 * @brief  加解密数据输出，16字节（128bit）分组数据输出
 * @param  初始向量
 * @retval None
 */

static void AES_DOUT_GroupRead_128BIT(uint8_t *dataOut)
{
    uint8_t       i;
    __IO uint32_t temp32;

    for (i = 0; i < 16; i = i + 4) {
        temp32         = FL_AES_ReadOutputData(AES);
        dataOut[0 + i] = temp32 >> 24;
        dataOut[1 + i] = temp32 >> 16;
        dataOut[2 + i] = temp32 >> 8;
        dataOut[3 + i] = temp32;
    }
}

/**
 * @brief  先设置16字节（128bit）加密数据，然后再读取16字节（128bit）解密数据
 * @param  加密数据
 * @param  解密数据
 * @retval None
 */
static FL_ErrorStatus AES_GroupWriteAndRead_128BIT(uint8_t *dataIn, uint8_t *dataOut)
{
    uint16_t i;

    /* 输入加解密数据 */
    AES_DIN_GroupWrite_128BIT(dataIn);

    /* 等待CCF标志置起,完成计算大约需要100个时钟周期 */
    for (i = 0; i < 500; i++) {
        if (FL_SET == FL_AES_IsActiveFlag_Complete(AES)) {
            break;
        }
    }

    /* 清CCF标志 */
    /* 读取前必须清除 */
    FL_AES_ClearFlag_Complete(AES);

    /* 读取结果 */
    AES_DOUT_GroupRead_128BIT(dataOut);

    if (i == 500) {
        return FL_FAIL;
    }

    else {
        return FL_PASS;
    }
}

/**
 * @brief  先设置16字节（128bit）加密数据，然后再读取16字节（128bit）解密数据
 * @param  加密数据
 * @param  解密数据
 * @retval None
 */
static FL_ErrorStatus AES_GroupWriteAndRead_128BIT_IVR(uint8_t *dataIn, uint8_t *dataOut)
{
    uint16_t i;

    /* 输入加解密数据 */
    AES_DIN_GroupWrite_128BIT(dataIn);

    /* 等待CCF标志置起,完成计算大约需要100个时钟周期 */
    for (i = 0; i < 500; i++) {
        if (FL_SET == FL_AES_IsActiveFlag_Complete(AES)) {
            break;
        }
    }

    /* 清CCF标志，读取前必须清除 */
    FL_AES_ClearFlag_Complete(AES);

    /* 读取结果 */
    AES_IVR_ReadEx(dataOut);

    if (i == 500) {
        return FL_FAIL;
    }

    else {
        return FL_PASS;
    }
}
